from typing import List
import asyncio
from dataprocess.elastic.EsSearcher import EsSearcher
from models.enums import EventMessageType
from Logger import Logger
from celery.result import AsyncResult
from llm.AnalyzeRelatedQuestion import AnalyzeRelatedQuestion
from llm.AnswerForUserQuestion import AnswerForUserQuestion
from llm.FinalAnswerToUser import FinalAnswerToUser
from llm.FinalAnswerByLLM import FinalAnswerByLLM
from utils.RedisQueue import RedisQueue
from utils.TextPaddingCleaner import TextPaddingCleaner
from utils.MessageEnrichUtils import MessageEnrichUtils
import time

logger = Logger.get_logger(__name__)

class RagEventHandler:
    """Event handler for task-related events"""
    
    @staticmethod
    async def handle(messages: dict):
        MAX_ROUND = 3
        try:
            question = messages["content"]
            _, all_question = await asyncio.to_thread(AnalyzeRelatedQuestion(user_requirement = question).run)
            # 补充信息，通过消息返回给前端
            messages = MessageEnrichUtils.enrich_related_questiones(messages, all_question)
            await asyncio.to_thread(RedisQueue().enqueue, EventMessageType.FRONT, messages)
            question_related_chunks =[]
            isFinal = False
            extra_info = None
            for round in range(MAX_ROUND):
                time.sleep(0.5)
                # 检索相关chunkes，返回前端
                related_chunks = await asyncio.to_thread(EsSearcher().search_by_embedding, all_question)
                messages = MessageEnrichUtils.enrich_related_chunks(messages, related_chunks)
                await asyncio.to_thread(RedisQueue().enqueue, EventMessageType.FRONT, messages)
                
                question_related_chunks.extend(related_chunks)
                # 合并chunkes，减小上下文
                cleared_question_related_chunks = TextPaddingCleaner.clean_text_padding(question_related_chunks)
                merge_for_chunks = TextPaddingCleaner.merge_for_context(cleared_question_related_chunks)
                # 尝试生成答案，并判断可信度
                isFinal, extra_info = await asyncio.to_thread(AnswerForUserQuestion(question, merge_for_chunks).run)
                messages["response"] = f'第{round+1}轮答案(可信度：{extra_info[1]}，注：超过90判定答案可靠并结束任务):\n{extra_info[0]}'
                await asyncio.to_thread(RedisQueue().enqueue, EventMessageType.FRONT, messages)
                if isFinal:
                    messages["response"] = '[DONE]'
                    await asyncio.to_thread(RedisQueue().enqueue, EventMessageType.FRONT, messages)
                    break
                all_question = extra_info[2]
            
            if not isFinal:
                time.sleep(2)
                _, final_answer = await asyncio.to_thread(FinalAnswerByLLM(question, []).run)
                messages["response"] = 'LLM生成答案：'+'\n'+final_answer[0]
                await asyncio.to_thread(RedisQueue().enqueue, EventMessageType.FRONT, messages)
                messages["response"] = '[DONE]'
                await asyncio.to_thread(RedisQueue().enqueue, EventMessageType.FRONT, messages)
            else:
                messages["response"] = '[DONE]'
                await asyncio.to_thread(RedisQueue().enqueue, EventMessageType.FRONT, messages) 
        except Exception as e:
            logger.error(f"Error handling rag event: {str(e)}")
